package edu.zju.gis.dbfg.model.tile;

import com.google.gson.Gson;
import com.google.gson.internal.LinkedTreeMap;
import edu.zju.gis.dbfg.model.PyramidConfig;
import edu.zju.gis.dbfg.model.base.Pipeline;
import edu.zju.gis.dbfg.model.base.PipelineBuilder;
import edu.zju.gis.dbfg.model.base.VectorTileBuilder;
import edu.zju.gis.dbfg.model.io.geojson.GeoJsonTileBuilder;
import edu.zju.gis.dbfg.model.io.pbf.MVTTileBuilder;
import edu.zju.gis.dbfg.model.util.GeoItem;
import edu.zju.gis.dbfg.model.util.TileID;
import edu.zju.gis.dbfg.model.util.TileUtil;
import edu.zju.gis.dbfg.model.util.ZLevelInfo;
import lombok.Getter;
import lombok.Setter;
import org.apache.http.HttpHost;
import org.apache.spark.api.java.JavaPairRDD;
import org.apache.spark.api.java.JavaRDD;
import org.apache.spark.api.java.JavaSparkContext;
import org.apache.spark.api.java.function.Function;
import org.apache.spark.api.java.function.PairFlatMapFunction;
import org.elasticsearch.action.search.ClearScrollRequest;
import org.elasticsearch.action.search.SearchRequest;
import org.elasticsearch.action.search.SearchResponse;
import org.elasticsearch.action.search.SearchScrollRequest;
import org.elasticsearch.client.RequestOptions;
import org.elasticsearch.client.RestClient;
import org.elasticsearch.client.RestHighLevelClient;
import org.elasticsearch.common.geo.ShapeRelation;
import org.elasticsearch.common.geo.builders.EnvelopeBuilder;
import org.elasticsearch.common.unit.TimeValue;
import org.elasticsearch.index.query.GeoShapeQueryBuilder;
import org.elasticsearch.search.Scroll;
import org.elasticsearch.search.SearchHit;
import org.elasticsearch.search.SearchHits;
import org.elasticsearch.search.builder.SearchSourceBuilder;
import org.geotools.geometry.jts.JTS;
import org.geotools.geometry.jts.ReferencedEnvelope;
import org.geotools.referencing.CRS;
import org.locationtech.jts.geom.*;
import org.locationtech.jts.geom.Polygon;
import org.opengis.referencing.FactoryException;
import org.opengis.referencing.crs.CoordinateReferenceSystem;
import org.opengis.referencing.operation.MathTransform;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import scala.Tuple2;

import java.awt.*;
import java.io.Serializable;
import java.util.*;
import java.util.List;

import static edu.zju.gis.dbfg.common.constant.TileConstant.SCREEN_TILE_BUFFER;
import static edu.zju.gis.dbfg.common.constant.TileConstant.SCREEN_TILE_SIZE;
import static edu.zju.gis.dbfg.model.tile.TileJob.OVERSAMPLE_FACTOR;
import static org.elasticsearch.index.query.QueryBuilders.geoShapeQuery;

@Getter
@Setter
public class MTileJob implements Serializable {

    private static final Logger logger = LoggerFactory.getLogger(MTileJob.class);

    private final static Integer SEARCH_SCROLL_SIZE = 2000;
    private final static long SEARCH_SCROLL_INTERVAL = 5L;
    private final static String SEARCH_ES_IP = "10.2.40.51, 10.2.40.52, 10.2.40.53, 10.2.40.54";
    private final static int SEARCH_ES_PORT = 9200;

    private PyramidConfig pyramidConfig;
    private String[] source;
    private int zmin;
    private int zmax;
    private CoordinateReferenceSystem sourceCrs;
    private CoordinateReferenceSystem targetCrs;
    private MathTransform stTransform;
    private MathTransform tsTransform;
    private String layerName;
    private String fileType;
    private String outdir;

    public MTileJob(PyramidConfig pyramidConfig, VectorTileClipperArgs mArgs, CoordinateReferenceSystem sourceCrs, CoordinateReferenceSystem targetCrs, String[] sources) throws FactoryException {
        this.pyramidConfig = pyramidConfig;
        this.zmin = mArgs.getZMin();
        this.zmax = mArgs.getZMax();
        this.fileType = mArgs.getFileType();
        this.outdir = mArgs.getOutDir();
        this.layerName = mArgs.getLayerName();
        this.source = sources;
        for (int i=0; i<this.source.length; i++) {
            this.source[i] = this.source[i].trim();
        }
        this.sourceCrs = sourceCrs;
        this.targetCrs = targetCrs;
        this.stTransform = CRS.findMathTransform(sourceCrs, targetCrs);
        this.tsTransform = CRS.findMathTransform(targetCrs, sourceCrs);
    }

    public void buildTile(JavaSparkContext jsc, TileID tileID) {
        List<TileID> taskTiles = new ArrayList<>();
        int z = tileID.getzLevel();
        int x = tileID.getX();
        int y = tileID.getY();
        int zTask = zmax;
        int xmin = x * (int) Math.pow(2, zTask-z);
        int xmax = ((x+1) * (int) Math.pow(2, zTask-z))-1;
        int ymin = y * (int) Math.pow(2, zTask-z);
        int ymax = ((y+1) * (int) Math.pow(2, zTask-z))-1;
        for (int i=xmin; i<=xmax; i++) {
            for (int j=ymin; j<=ymax; j++) {
                TileID l = new TileID();
                l.setX(i);
                l.setY(j);
                l.setzLevel(zTask);
                taskTiles.add(l);
            }
        }

        JavaPairRDD<TileID, GeoItem> tileFeatures = jsc.parallelize(taskTiles, taskTiles.size()).flatMapToPair(new PairFlatMapFunction<TileID, TileID, GeoItem>() {
            @Override
            public Iterator<Tuple2<TileID, GeoItem>> call(TileID tileID) throws Exception {
                List<GeoItem> features = readFeatures(tileID);
                List<Tuple2<TileID, GeoItem>> result = new ArrayList<>();
                if (features == null || features.size() == 0) return result.iterator();
                Envelope tileEnvelope = TileUtil.createTileBox(tileID, pyramidConfig);
                Pipeline prePipeline = getPipeline(targetCrs, SCREEN_TILE_BUFFER, tileEnvelope, true, false, false);
                Pipeline pipeline = getPipeline(targetCrs, SCREEN_TILE_BUFFER, tileEnvelope, false, true, true);

                /**
                 * 写出当前层级
                 */
                VectorTileBuilder vectorTileBuilder;
                if (fileType.equals("geojson")) {
                    vectorTileBuilder = new GeoJsonTileBuilder(outdir, tileID, pyramidConfig);
                } else {
                    vectorTileBuilder = new MVTTileBuilder(outdir, tileID, tileEnvelope, layerName);
                }

                for (GeoItem gi: features) {
                    Geometry g = gi.getGeometry();
                    g = prePipeline.execute(g);
                    if (!g.isEmpty()) {
                        g = pipeline.execute(g);
                        if (!g.isEmpty()) {
                            vectorTileBuilder.addFeature(layerName, gi.getFid(), gi.getFid(), g, new HashMap<>());
                        }
                    }
                }

                vectorTileBuilder.build();

                /**
                 * 计算上一层级并写出
                 */
                TileID upperTileID = new TileID();
                upperTileID.setzLevel(tileID.getzLevel()-1);
                upperTileID.setX(tileID.getX()/2);
                upperTileID.setY(tileID.getY()/2);
                Envelope upperLevelEnvelope = TileUtil.createTileBox(upperTileID, pyramidConfig);

                Pipeline upperPipeline = getPipeline(targetCrs, SCREEN_TILE_BUFFER, upperLevelEnvelope, false, false, true);
                for (GeoItem gi: features) {
                    Geometry g = gi.getGeometry();
                    g = upperPipeline.execute(g);
                    if (!g.isEmpty()) {
                        g = onlyPolygon(g);
                        if (g != null) {
                            gi.setGeometry(g);
                            result.add(new Tuple2<>(upperTileID, gi));
                        }
                    }
                }

                return result.iterator();
            }
        });

        for (int i=zmax-1; i>=zmin; i--) {
            tileFeatures = tileFeatures.groupByKey().flatMapToPair(new TileBuilder());
        }

        logger.info(String.format("Job Tile %s succeeded: %d", tileID.toString(), tileFeatures.count()));
    }

    public List<GeoItem> readFeatures(TileID tileID) throws Exception {

        String[] ips = SEARCH_ES_IP.split(",");
        HttpHost[] hosts = new HttpHost[ips.length];
        for (int m = 0; m < ips.length; m++) {
            hosts[m] = new HttpHost(ips[m].trim(), SEARCH_ES_PORT);
        }
        RestHighLevelClient client = new RestHighLevelClient(
                RestClient.builder(hosts));

        List<GeoItem> result = new ArrayList<>();

        /**
         * 查询 es 获取当前格网内部的数据
         */
        try {
            List<Tuple2<String, String>> tuples = new ArrayList<>();
            Envelope tile = TileUtil.createTileBox(tileID, pyramidConfig);

            tile = JTS.transform(tile, tsTransform);

            SearchSourceBuilder sourceBuilder = new SearchSourceBuilder();
            GeoShapeQueryBuilder qb = geoShapeQuery(
                    "shape",
                    new EnvelopeBuilder(new Coordinate(tile.getMinX(), tile.getMaxY()), new Coordinate(tile.getMaxX(), tile.getMinY())));
            qb.relation(ShapeRelation.INTERSECTS);
            sourceBuilder.query(qb);
            sourceBuilder.size(SEARCH_SCROLL_SIZE);
            // first search request to get one batch of data
            SearchRequest searchRequest = new SearchRequest();
            searchRequest.indices(source);
            searchRequest.source(sourceBuilder);
            searchRequest.scroll(TimeValue.timeValueMinutes(SEARCH_SCROLL_INTERVAL));
            SearchResponse searchResponse = client.search(searchRequest, RequestOptions.DEFAULT);
            SearchHits hits = searchResponse.getHits();
            SearchHit[] searchHits = hits.getHits();
            for (SearchHit hit : searchHits) {
                tuples.add(new Tuple2<>(hit.getId(), hit.getSourceAsString()));
            }

            // scroll request to get all left relevant data
            String scrollId = searchResponse.getScrollId();
            final Scroll scroll = new Scroll(TimeValue.timeValueMinutes(SEARCH_SCROLL_INTERVAL));
            while (searchHits != null && searchHits.length > 0) {
                SearchScrollRequest scrollRequest = new SearchScrollRequest(scrollId);
                scrollRequest.scroll(scroll);
                searchResponse = client.scroll(scrollRequest, RequestOptions.DEFAULT);
                scrollId = searchResponse.getScrollId();
                searchHits = searchResponse.getHits().getHits();
                for (SearchHit hit : searchHits) {
                    tuples.add(new Tuple2<>(hit.getId(), hit.getSourceAsString()));
                }
            }

            // clear scroll id
            ClearScrollRequest clearScrollRequest = new ClearScrollRequest();
            clearScrollRequest.addScrollId(scrollId);

            client.clearScroll(clearScrollRequest, RequestOptions.DEFAULT);
            client.close();

            if (tuples.size() == 0) {
                return result;
            }

            logger.info(String.format("TileID %s, Tuples count: %d", tileID.toString(), tuples.size()));

            /**
             * 对于获取到的数据,转为GeoItem
             */
            for (Tuple2<String, String> f : tuples) {
                Gson gson = new Gson();
                Map<String, Object> inMap = gson.fromJson(f._2, new HashMap<String, Object>().getClass());
                inMap.remove("id");
                if (!inMap.containsKey("shape")) {
                    logger.error("Unvalid geometry item for id " + f._1 + "with source " + source + ", must have field named with 'shape'");
                    continue;
                }
                Map<String, Object> geoMap = (LinkedTreeMap) inMap.get("shape");
                String type = (String) geoMap.get("type");
                ArrayList t = ((ArrayList) geoMap.get("coordinates"));
                GeometryFactory gf = new GeometryFactory();
                for (int i = 0; i < t.size(); i++) {
                    List<List<ArrayList>> coordinates = new ArrayList<>();
                    ArrayList<ArrayList> m = (ArrayList<ArrayList>) t.get(i);
                    if (!type.contains("Multi")) {
                        coordinates.add(m);
                    } else {
                        for (int x = 0; x < m.size(); x++) {
                            coordinates.add(m.get(x));
                        }
                    }
                    for (List<ArrayList> coordinate : coordinates) {
                        Geometry geometry;
                        if (type.equals("Polygon") || type.equals("MultiPolygon")) {
                            Coordinate[] cs = coordinate.stream().map((p) -> new Coordinate((double) (p.get(0)), (double) (p.get(1)))).toArray(Coordinate[]::new);
                            geometry = gf.createPolygon(cs);
                        } else if (type.equals("LineString") || type.equals("MultiLineString")) {
                            Coordinate[] cs = coordinate.stream().map((p) -> new Coordinate((double) (p.get(0)), (double) (p.get(1)))).toArray(Coordinate[]::new);
                            geometry = gf.createLineString(cs);
                        } else {
                            // type == Point
                            Coordinate c = new Coordinate((double) (coordinate.get(0).get(0)), (double) (coordinate.get(0).get(1)));
                            geometry = gf.createPoint(c);
                        }
                        GeoItem geoItem = new GeoItem();
                        geoItem.setFid(f._1);
                        inMap.remove("shape");
                        geometry = JTS.transform(geometry, stTransform);
                        if (!geometry.isEmpty()) {
                            geometry = onlyPolygon(geometry);
                            if (geometry != null) {
                                geoItem.setGeometry(geometry);
                                // geoItem.setProp(inMap);
                                result.add(geoItem);
                            }
                        }
                    }
                }
            }
            logger.info(String.format("TileID %s, Features count: %d", tileID.toString(), result.size()));
            return result;
        } catch (Exception e) {
            logger.error(String.format("Tile %s fail: %s", tileID.toString(), e.getMessage()));
        } finally {
            client.close();
        }
        return null;
    }

    /**
     * 构建瓦片中每个像素处理的Pipeline
     * 默认目标crs与源crs一致
     */
    public Pipeline getPipeline(CoordinateReferenceSystem sourceCrs, int buffer, Envelope referenceEnvelop, boolean preprocess, boolean isClip, boolean isSimplify) throws Exception {
        CoordinateReferenceSystem targetCrs = sourceCrs;
        int mapWidth = SCREEN_TILE_SIZE;
        int mapHeight = SCREEN_TILE_SIZE;
        int overSampleFactor = OVERSAMPLE_FACTOR;
        // rectangle of the image: width and height
        Rectangle paintArea = new Rectangle(mapWidth, mapHeight);

        // bounding box - in final map (target) CRS
        ReferencedEnvelope renderingArea = new ReferencedEnvelope(referenceEnvelop.getMinX(), referenceEnvelop.getMaxX(), referenceEnvelop.getMinY(), referenceEnvelop.getMaxY(), targetCrs);
        Pipeline pipeline = null;

        try {
            final PipelineBuilder builder =
                    PipelineBuilder.newBuilder(
                            renderingArea, paintArea, sourceCrs, overSampleFactor, buffer);

            if (preprocess){
                pipeline =
                        builder.preprocess()
                                .collapseCollections()
//                                .precision(targetCrs)
                                .build();
            } else  {
                pipeline = builder
                        .clip(isClip, false)
                        .simplify(isSimplify)
                        .collapseCollections()
//                        .precision(targetCrs)
                        .build();
            }
        } catch (FactoryException e) {
            e.printStackTrace();
        }

        return pipeline;
    }


    public Geometry onlyPolygon(Geometry result) {
        if ((result instanceof org.locationtech.jts.geom.Polygon) || (result instanceof MultiPolygon)) {
            return result;
        }
        List polys = org.locationtech.jts.geom.util.PolygonExtracter.getPolygons(result);
        if (polys.size() == 0) {
            return null;
        }
        if (polys.size() == 1) {
            return (org.locationtech.jts.geom.Polygon) polys.get(0);
        }
        return new MultiPolygon(
                (org.locationtech.jts.geom.Polygon[]) polys.toArray(new Polygon[polys.size()]), result.getFactory());
    }

    /**
     * 将裁切并抽稀的图斑合并
     * 如果直接合并,会有两个图斑不存在相交点的报错
     * @param g1
     * @param g2
     * @param buffer
     * @return
     */
    public Geometry union(Geometry g1, Geometry g2, double buffer, int depth, int maxDepth) {

        Geometry bg1;
        Geometry bg2;

        if (depth == 0) {
            bg1 = g1;
            bg2 = g2;
        } else {
            bg1 = g1.buffer(buffer);
            bg2 = g2.buffer(buffer);
        }

        try {
            return bg1.union(bg2);
        } catch (TopologyException e) {
            depth ++;
            if (depth > maxDepth) {
                return null;
            }
            return union(g1, g2, buffer*2, depth, maxDepth);
        }
    }

    private class TileBuilder implements PairFlatMapFunction<Tuple2<TileID, Iterable<GeoItem>>, TileID, GeoItem> {
        @Override
        public Iterator<Tuple2<TileID, GeoItem>> call(Tuple2<TileID, Iterable<GeoItem>> tileIDIterableTuple2) throws Exception {

            List<Tuple2<TileID, GeoItem>> result = new ArrayList<>();

            TileID tileID = tileIDIterableTuple2._1;

            if (!tileIDIterableTuple2._2.iterator().hasNext()) {
                logger.warn(String.format("tile[%d, %d, %d] abort with total count [%d]", tileID.getzLevel(), tileID.getX(), tileID.getY(), 0));
                return result.iterator();
            }

            Iterator<GeoItem> geoItemIterator = tileIDIterableTuple2._2.iterator();
            Envelope tileEnvelope = TileUtil.createTileBox(tileID, pyramidConfig);
            Pipeline preProcessPipeline = getPipeline(targetCrs, SCREEN_TILE_BUFFER, tileEnvelope, true, false, false);

            VectorTileBuilder vectorTileBuilder;
            if (fileType.equals("geojson")) {
                vectorTileBuilder = new GeoJsonTileBuilder(outdir, tileID, pyramidConfig);
            } else {
                vectorTileBuilder = new MVTTileBuilder(outdir, tileID, tileEnvelope, layerName);
            }

            TileID upperTileID = new TileID();
            upperTileID.setzLevel(tileID.getzLevel()-1);
            upperTileID.setX(tileID.getX()/2);
            upperTileID.setY(tileID.getY()/2);

            /**
             * 将之前被分割的图形union回来
             */
            Map<String, GeoItem> unionedGeoItem = new HashMap<>();

            while(geoItemIterator.hasNext()) {
                GeoItem geoItem = geoItemIterator.next();
                if (!unionedGeoItem.containsKey(geoItem.getFid())) {
                    unionedGeoItem.put(geoItem.getFid(), geoItem);
                } else {
                    Geometry g = onlyPolygon(geoItem.getGeometry());
                    Geometry og = onlyPolygon(unionedGeoItem.get(geoItem.getFid()).getGeometry());
                    if (og == null) continue;
                    Geometry ug = union(g, og, Math.max(tileEnvelope.getWidth()/(16*SCREEN_TILE_SIZE), tileEnvelope.getHeight()/(16*SCREEN_TILE_SIZE)), 0, 2);
                    if (ug == null) {
                        continue;
                    }
                    geoItem.setGeometry(ug);
                    // Map<String, Object> prop = geoItem.getProp();
                    // prop.put("k", (Double)prop.getOrDefault("k", 0d) + (Double)unionedGeoItem.get(geoItem.getFid()).getProp().getOrDefault("k", 0d));
                    // geoItem.setProp(prop);
                    unionedGeoItem.put(geoItem.getFid(), geoItem);
                }
            }

            List<Tuple2<TileID, GeoItem>> temp = new ArrayList<>();

            Set<String> rKeys = unionedGeoItem.keySet();
            for (String key: rKeys) {
//                String key = keys.next();
                GeoItem geoItem = unionedGeoItem.get(key);
                // Geometry clippedGeom = clipPipeline.execute(geoItem.getGeometry());
                Geometry finalGeom = preProcessPipeline.execute(geoItem.getGeometry());
                if (!finalGeom.isEmpty()) {
                    finalGeom = onlyPolygon(finalGeom);
                    if (finalGeom == null) continue;
                    vectorTileBuilder.addFeature(layerName, geoItem.getFid(), geoItem.getFid(),
                            finalGeom, new HashMap<>());
                    geoItem.setGeometry(finalGeom);
                    temp.add(new Tuple2<>(upperTileID, geoItem));
                }
            }

            vectorTileBuilder.build();
            logger.info(String.format("tile[%d, %d, %d] success with total count [%d]", tileID.getzLevel(), tileID.getX(), tileID.getY(), temp.size()));

            /**
             * 处理数据用于上一层计算，这时就不需要clip了，只需要进行抽稀即可
             */
            if (tileID.getzLevel() > zmin) {
                Envelope upperLevelEnvelope = TileUtil.createTileBox(upperTileID, pyramidConfig);
                Pipeline upperLevelSimplifyPipeline = getPipeline(targetCrs, SCREEN_TILE_BUFFER, upperLevelEnvelope, false, false, true);
                for (int i=0; i<temp.size(); i++) {
                    GeoItem gi = temp.get(i)._2;
                    TileID ti = temp.get(i)._1;
                    Geometry g = upperLevelSimplifyPipeline.execute(gi.getGeometry());
                    if (g.isEmpty()) {
                        continue;
                    }
                    g = onlyPolygon(g);
                    if (g == null) {
                        continue;
                    }
                    gi.setGeometry(g);
                    result.add(new Tuple2<>(ti, gi));
                }
            } else {
                result.add(new Tuple2<>(null, null));
            }

            return result.iterator();
        }

    }
}
